package zhao.exmaple;


import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.streaming.api.datastream.ConnectedStreams;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.CoMapFunction;
import org.apache.flink.streaming.api.functions.co.CoProcessFunction;
import org.apache.flink.util.Collector;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class ConnectKeyByDemo {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(2);

        DataStreamSource<Tuple2<Integer, String>> source1 = env.fromElements(
                Tuple2.of(1, "a1"),
                Tuple2.of(1, "a1"),
                Tuple2.of(2, "a2"),
                Tuple2.of(3, "a3")

        );
        DataStreamSource<Tuple3<Integer, String, Integer>> sourc2 = env.fromElements(
                Tuple3.of(1, "aa1", 1),
                Tuple3.of(1, "aa2", 2),
                Tuple3.of(2, "bb", 1),
                Tuple3.of(3, "cc", 1)
        );

        ConnectedStreams<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>> connect = source1.connect(sourc2);

        // keyBy 對 key進行分組，在并行執行時分組操作
        ConnectedStreams<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>> tuple2Tuple3ConnectedStreams = connect.keyBy(s1 -> s1.f0, s2 -> s2.f0);

        SingleOutputStreamOperator process = tuple2Tuple3ConnectedStreams.process(new CoProcessFunction<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>, String>() {
        // SingleOutputStreamOperator process = connect.process(new CoProcessFunction<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>, String>() {
            Map<Integer, List<Tuple2<Integer, String>>> sache1 = new HashMap<>();
            Map<Integer, List<Tuple3<Integer, String, Integer>>> sache2 = new HashMap<>();

            @Override
            public void processElement1(Tuple2<Integer, String> value, CoProcessFunction<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>, String>.Context ctx, Collector<String> out) throws Exception {
                Integer id = value.f0;
                if (!sache1.containsKey(id)) {
                    List<Tuple2<Integer, String>> list = new ArrayList<>();
                    list.add(value);
                    sache1.put(id, list);
                } else {
                    sache1.get(id).add(value);
                }
                if (sache2.containsKey(id)) {
                    for (Tuple3<Integer, String, Integer> tuple3 : sache2.get(id)) {
                        out.collect("s1" + value + "=====>  s2" + tuple3);
                    }
                }

            }

            @Override
            public void processElement2(Tuple3<Integer, String, Integer> value, CoProcessFunction<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>, String>.Context ctx, Collector<String> out) throws Exception {
                Integer id = value.f0;

                if (!sache2.containsKey(id)) {
                    List<Tuple3<Integer, String, Integer>> list = new ArrayList<>();
                    list.add(value);
                    sache2.put(id, list);
                } else {
                    sache2.get(id).add(value);
                }
                if (sache1.containsKey(id)) {
                    for (Tuple2<Integer, String> tuple2 : sache1.get(id)) {
                        out.collect("s1" + tuple2 + "=====>  s2" + value);
                    }
                }
            }

        });

        process.print();


        env.execute();

    }
}
